#!/usr/bin/python
#
# faqemu - programmable script with an interface similar to that of qemu.
#          should help debugging vdsm

import sys
import os
import getopt
import time
import code
import socket
import threading
import SocketServer
import pickle
import traceback
import StringIO

import qGuestAgentFake

STARTUP_SLEEP = 2

stopped = threading.Event()

guestAgent = None
drives = []
events = []

running = True
mstate = ''

def usage():
    print 'faqemu: use normal qemu parameters'
    sys.exit(1)

def appendDrive(drvarg):
    d = {}
    for elem in drvarg.split(','):
        kv = elem.split('=')
        if len(kv) == 2:
            d[kv[0]] = kv[1]
    drives.append(d)

def driveLookup(media):
    for d in drives:
        if d.get('media') == media:
            return d
    return None

MIGRATION_MAGIC = '\x71\x84'
def migrateTo(dst, port):
    global running, mstate
    mstate = 'ongoing'
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((dst, port))
        try:
            sock.send(MIGRATION_MAGIC)
            rmagic = sock.recv(len(MIGRATION_MAGIC))
            if rmagic != MIGRATION_MAGIC:
                raise RuntimeError('Unexpected magic %s' % repr(rmagic))
            else:
                pickle.dump({'version': 1}, sock.makefile())
                mstate = 'completed'
                running = False
                return True
        finally:
            sock.close()
    except:
        print traceback.format_exc()
        mstate = 'failed'
        return False

def migrateFrom(port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(('0', port))
    sock.listen(5)
    csock, caddr = sock.accept()
    csock.send(MIGRATION_MAGIC)
    rmagic = csock.recv(len(MIGRATION_MAGIC))
    try:
        if rmagic != MIGRATION_MAGIC:
            print 'Unexpected magic %s' % repr(rmagic)
        else:
            return pickle.load(csock.makefile())
    finally:
        csock.close()
        sock.close()

def run():
    monitorSocketName = hcSocketName = pidFileName = interpretSocketName = None
    guestLogFileName = None
    uuid = None
    migrationDest = None

    # ugly hack to support single-dash long options
    argv = sys.argv[1:]
    for i in range(0, len(argv)):
        if argv[i][0] == '-' and argv[i][1] != '-' and len(argv[i]) > 2:
            argv[i] = '-' + argv[i]

    opts, args = getopt.getopt(argv, 'm:L:k:SM:', ['help', 'monitor=', 'net=',
        'hda=', 'pidfile=', 'vmchannel=', 'cpu-vendor=', 'no-kvm',
        'localtime', 'startdate=', 'usbdevice=', 'vnc=', 'fda=', 'soundhw=',
        'spice=', 'uuid=', 'no-kvm-pit-reinjection', 'rtc-td-hack', 'name=',
        'smp=', 'boot=', 'net=', 'notify=', 'balloon=', 'smbios=', 'qxl=',
        'spice=', 'soundhw=', 'drive=', 'cpu=', 'incoming=' ])

    for o, v in opts:
        if o == '--help':
            usage()
        elif o == '--monitor':
            monitorSocketName = v[5:-7]
        elif o == '--vmchannel':
            hcSocketName = v[13:-7]
        elif o == '--pidfile':
            pidFileName = v
        elif o == '--uuid':
            uuid = v
        elif o == '--smbios':
            for elem in v.split(','):
                kv = elem.split('=')
                if len(kv) == 2 and kv[0] == 'uuid':
                    uuid = kv[1]
        elif o == '--drive':
            appendDrive(v)
        elif o == '--incoming':
            migrationDest = v
        elif o == '-S':
            running = False
        elif o == '-M' and v == '?':
            print "pc         Standard PC (default)\n" +\
                  "rhel5.4.0  RHEL 5.4.0 PC\n" +\
                  "rhel5.5.0  RHEL 5.5.0 PC"
            return 0

    if uuid:
        interpretSocketName = '/var/vdsm/%s.interpreter.socket' % uuid
        guestLogFileName = '/var/vdsm/%s.guest.log' % uuid

    if pidFileName:
        file(pidFileName, 'w').write(str(os.getpid()))

    time.sleep(STARTUP_SLEEP)

    if monitorSocketName:
        removef(monitorSocketName)
        mon = SocketServer.UnixStreamServer(monitorSocketName, MonitorHandler)
        mt = threading.Thread(target=mon.serve_forever, name='Monitor')
        mt.setDaemon(True)
        mt.start()

    if interpretSocketName:
        removef(interpretSocketName)
        inter = SocketServer.UnixStreamServer(interpretSocketName, InterpreterHandler)
        os.chmod(interpretSocketName, 0777)
        it = threading.Thread(target=inter.serve_forever, name='Interpreter')
        it.setDaemon(True)
        it.start()

    if migrationDest:
        method, dst = migrationDest.split(':', 1)
        if method == 'tcp':
            dst, port = dst.split(':', 1)
            port = int(port)
            migdata = migrateFrom(port)
            if not migdata:
                sys.exit(1)
            print 'migration finished successfully'
            events.append('# migration: migration process finished')


    global guestAgent
    guestAgent = qGuestAgentFake.FakeVdsAgent(guestLogFileName)
    guestAgent.run()

    if hcSocketName:
        removef(hcSocketName)
        hc = SocketServer.UnixStreamServer(hcSocketName, AgentHandler)
        hct = threading.Thread(target=hc.serve_forever, name='Agent')
        hct.setDaemon(True)
        hct.start()

    stopped.wait()
    print 'bye!'

def removef(f):
    try:
        os.remove(f)
    except:
        pass


class InterpreterHandler(SocketServer.StreamRequestHandler):
    _interpreter = code.InteractiveConsole({'events':events})

    def handle(self):
        while True:
            needMoreInput = False
            result = ""

            line = self.rfile.readline()
            # In case someone did something stupid and stuck the interpreter
            # reset will free the buffer.
            if line.strip() == "#RESET":
                buff = "\n".join(self._interpreter.buffer)
                self._interpreter.resetbuffer()
                result = buff
                needMoreInput = False
            else:
                needMoreInput, result = self._pushAndCapture(line)

            if not needMoreInput:
                if result == None:
                    result = ""
                size = hex(len(result))[2:].zfill(4)
                self.wfile.write(size + result)

    def _pushAndCapture(self, line):
        """
        Push a line to the inner interpreter and capture stdout and stderr.

        :returns: a tuple (need more info, stdout and stderr)

        need more input states whether the line was (or ended) a complete python expression.
        (ie. `if a == b:` is not a complete statement and more lines are needed)
        """
        # XXX playing with the global stdout is a bad bad idea.
        out, err = sys.stdout, sys.stderr
        try:
            sys.stdout = sys.stderr = StringIO.StringIO()
            needMoreInput = self._interpreter.push(line)
            if not needMoreInput:
                return (needMoreInput, sys.stdout.getvalue())

            return (needMoreInput, None)
        finally:
            sys.stdout, sys.stderr = out, err


class MonitorHandler(SocketServer.StreamRequestHandler):
    def _writeEvents(self):
        try:
            while True:
                self.wfile.write('%s\r\n' % events.pop())
        except:
            pass

    def _writeLoop(self):
        self._writeEvents()
        threading.Timer(1.0, self._writeLoop).start()

    def handle(self):
        self.wfile.write('FAQEMU monitor\n')
        threading.Timer(1.0, self._writeLoop).start()
        while True:
            try:
                self.wfile.write('(qemu) ')
                cmd = self.rfile.readline().strip()
                if not cmd: break
                if cmd == 'quit':
                    stopped.set()
                    break
                elif cmd == 'info block':
                    ret = ''
                    hda = driveLookup('disk')
                    if hda:
                        ret += "ide0-hd0: type=hd removable=0 file=%s ro=0 drv=qcow2 encrypted=0\r\n" % hda['file']
                    else:
                        ret += "ide1-hd0: type=hd removable=1 locked=0 [not inserted]\r\n"
                    cdrom = driveLookup('cdrom')
                    if cdrom:
                        ret += "ide0-cd0: type=cdrom removable=0 file=%s ro=0 drv=qcow2 encrypted=0\r\n" % cdrom['file']
                    else:
                        ret += "ide1-cd0: type=cdrom removable=1 locked=0 [not inserted]\r\n"
                    ret += """floppy0: type=floppy removable=1 locked=0 [not inserted]
sd0: type=floppy removable=1 locked=0 [not inserted]\r\n"""
                    self.wfile.write(ret)
                elif cmd == 'info blockstats':
                    ret = """ide0-hd0: rd_bytes=0 wr_bytes=0 rd_operations=0 wr_operations=0\r
ide1-cd0: rd_bytes=0 wr_bytes=0 rd_operations=0 wr_operations=0\r
floppy0: rd_bytes=0 wr_bytes=0 rd_operations=0 wr_operations=0\r
sd0: rd_bytes=0 wr_bytes=0 rd_operations=0 wr_operations=0\r
"""
                    self.wfile.write(ret)
                elif cmd.startswith('migrate -d'):
                    dst = cmd.split()[2]
                    method, dst = dst.split(':', 1)
                    if method == 'tcp':
                        dst, port = dst.split(':', 1)
                        port = port.split(',', 1)[0]
                        port = int(port)
                        t = threading.Thread(target=migrateTo, args=(dst, port))
                        t.setDaemon(True)
                        t.start()
                    else:
                        print 'unsupported cmd %s' % cmd
                    self.wfile.write('\r\n')
                elif cmd == 'info status':
                    if running:
                        self.wfile.write('VM status: running\r\n')
                    else:
                        self.wfile.write('VM status: paused\r\n')
                elif cmd == 'info migrate':
                    if mstate:
                        self.wfile.write('Migration status: %s\r\n' % mstate)
                elif cmd == 'info kvm':
                    self.wfile.write('kvm support: enabled\r\n')
                else:
                    self.wfile.write('received command %s\r\n' % cmd)
            except Exception, e:
                self.wfile.write(str(e.__class__) + str(e) + '\r\n')
                print >> sys.stderr, traceback.format_exc()

class AgentHandler(SocketServer.StreamRequestHandler):
    def handle(self):
        guestAgent.hc.setSocket(self.request)
        guestAgent.listenThread.join()

run()
